import random

import gym
import numpy as np
import pygame
from cpgames.core import PingpongGame
from cpgames.core.games.pingpong.modules import Racket, Ball
from cpgames.core.utils import QuitGame


class NoSound:
    def play(self):
        pass


class AIBall(Ball):

    def reset(self):
        super().reset()
        self.speed = random.randint(1, 6)

    def move(self, ball, racket_left, racket_right, hit_sound, goal_sound):
        scores = super(AIBall, self).move(ball, racket_left, racket_right, hit_sound, goal_sound)
        # self.speed = min(self.speed, 10)
        return scores


class AIPingPongGame(PingpongGame):

    def AIGamingInterface(self, screen, resource_loader, cfg):
        # 初始化
        # hit_sound = resource_loader.sounds['hit']
        # goal_sound = resource_loader.sounds['goal']
        hit_sound = NoSound()
        goal_sound = NoSound()
        font = resource_loader.fonts['default50']
        # resource_loader.playbgm()
        # 开始界面
        # game_mode = 1
        # 游戏主循环
        # --左边球拍(ws控制, 仅双人模式时可控制)
        # score_left = 0
        racket_left = Racket(cfg.IMAGE_PATHS_DICT['racket'], 'LEFT', cfg)
        # --右边球拍(↑↓控制)
        # score_right = 0
        racket_right = Racket(cfg.IMAGE_PATHS_DICT['racket'], 'RIGHT', cfg)
        # racket_right.speed = 10
        # --球
        ball = AIBall(cfg.IMAGE_PATHS_DICT['ball'], cfg)

        return racket_left, racket_right, ball, hit_sound, goal_sound, font


class PingPongEnv(gym.Env):
    def __init__(self, game):
        self.game = AIPingPongGame()
        self.score_left = 0
        self.score_right = 0
        self.clock = pygame.time.Clock()
        self.observation_space = gym.spaces.Discrete(7)
        self.action_space = gym.spaces.Discrete(3)
        self.last_action = 2

    def on_hot_obs(self):
        # 将球的x坐标转换为one-hot编码
        ball_x = np.zeros(500)
        ball_x[self.ball.rect.x-1] = 1
        # 将球的y坐标转换为one-hot编码
        ball_y = np.zeros(500)
        ball_y[self.ball.rect.y-1] = -1
        # 相加得到一个长度为500的向量
        ball = ball_x + ball_y
        # 将球的速度方向为one-hot编码
        if self.ball.direction_x < 0 and self.ball.direction_y < 0:
            ball_direction = [self.ball.speed / 10, 0, 0, 0]
        elif self.ball.direction_x < 0 and self.ball.direction_y > 0:
            ball_direction = [0, self.ball.speed / 10, 0, 0]
        elif self.ball.direction_x > 0 and self.ball.direction_y < 0:
            ball_direction = [0, 0, self.ball.speed / 10, 0]
        else:
            ball_direction = [0, 0, 0, self.ball.speed / 10]
        # 将球拍的y坐标转换为one-hot编码
        racket_y = np.zeros(500)
        racket_y[self.racket_right.rect.y] = 1
        # 链接
        return np.concatenate((ball, ball_direction, racket_y))

    def reset(self, *args, **kwargs):
        racket_left, racket_right, ball, hit_sound, goal_sound, font = self.game.AIGamingInterface(
            self.game.screen,
            self.game.resource_loader,
            self.game.cfg
        )
        self.racket_left = racket_left
        self.racket_right = racket_right
        self.score_right = 0
        self.score_left = 0
        self.ball = ball
        self.hit_sound = hit_sound
        self.goal_sound = goal_sound
        self.last_action = 2
        self.font = font

        # # 返回状态, 用于初始化 [球速, 球位置, 球方向, 左拍位置, 右拍位置]
        # return self.ball.speed / 10, \
        #        (self.ball.rect.centerx - 250) / 250, \
        #        (self.ball.rect.centery - 250) / 250, \
        #        self.ball.direction_x, self.ball.direction_y, \
        #        (self.racket_right.rect.centery - 250) / 250, \
        #        (self.racket_left.rect.centery - 250) / 250
        return self.on_hot_obs()

    def step(self, action):
        # action = 0, 1, 2分别代表上，下，不动
        if action == 0:
            self.racket_right.move('UP')
        elif action == 1:
            self.racket_right.move('DOWN')
        elif action == 2:
            pass
        # 球运动
        scores = self.ball.move(self.ball, self.racket_left, self.racket_right, self.hit_sound, self.goal_sound)
        self.racket_left.automove(self.ball)
        self.score_left += scores[0]
        self.score_right += scores[1]
        done = False
        if self.score_left == 11 or self.score_right == 11:
            done = True
            # reward = -abs(self.racket_right.rect.centery - self.ball.rect.centery) * 0.1
        # 掉球惩罚
        racket_height = self.racket_right.rect.bottom - self.racket_right.rect.top
        if scores[0] > 0:
            # 距离中心越远,分数越低
            reward = -(abs(self.racket_right.rect.centery - self.ball.rect.centery) / racket_height) * 2
        elif scores[1] > 0:
            reward = 1
        # 打中奖励
        elif pygame.sprite.collide_rect(self.ball, self.racket_right):
            # 计算球拍高度
            # 距离中心越近，分数越高
            if abs(self.ball.rect.centery - self.racket_right.rect.centery) < racket_height / 4:
                reward = (racket_height / 2 - abs(
                    self.ball.rect.centery - self.racket_right.rect.centery)) / racket_height / 5
            else:
                reward = 0
        else:
            reward = 0
        # if abs(self.racket_right.rect.centery - self.ball.rect.centery) > 20:
        #     reward += -0.001
        # if action == self.last_action:
        #     reward += 0.00001
        # else:
        #     reward += -0.00001
        self.last_action = action
        # 返回状态, 奖励, 是否结束, 其他信息
        # return (self.ball.speed / 10,
        #         (self.ball.rect.centerx - 250) / 250,
        #         (self.ball.rect.centery - 250) / 250,
        #         self.ball.direction_x, self.ball.direction_y,
        #         (self.racket_right.rect.centery - 250) / 250,
        #         (self.racket_left.rect.centery - 250) / 250), reward, done, {}
        return self.on_hot_obs(), reward, done, {}

    def render(self, mode='human'):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                QuitGame()
        self.game.screen.fill((41, 36, 33))
        # 玩家操作
        pressed_keys = pygame.key.get_pressed()
        if pressed_keys[pygame.K_UP]:
            self.racket_right.move('UP')
        elif pressed_keys[pygame.K_DOWN]:
            self.racket_right.move('DOWN')
        self.racket_left.automove(self.ball)
        # 球运动
        scores = self.ball.move(self.ball, self.racket_left, self.racket_right, self.hit_sound, self.goal_sound)
        self.score_left += scores[0]
        self.score_right += scores[1]
        # 显示
        # --分隔线
        pygame.draw.rect(self.game.screen, self.game.cfg.WHITE, (247, 0, 6, 500))
        # --球
        self.ball.draw(self.game.screen)
        # --拍
        self.racket_left.draw(self.game.screen)
        self.racket_right.draw(self.game.screen)
        # --得分
        self.game.screen.blit(self.font.render(str(self.score_left), False, self.game.cfg.WHITE), (150, 10))
        self.game.screen.blit(self.font.render(str(self.score_right), False, self.game.cfg.WHITE), (300, 10))
        pygame.display.update()
        self.clock.tick(self.game.cfg.FPS_GAMING)


if __name__ == '__main__':
    pingpone_env = PingPongEnv(AIPingPongGame())
    # pingpone_env.game.run()
    while True:
        state = pingpone_env.reset()
        done = False
        while not done:
            action = random.randint(0, 3)
            state, reward, done, _ = pingpone_env.step(action)
            pingpone_env.render()
            print(state, reward, done)
